#include "lanenet_inference.hpp"
#include <iostream>

LanenetInference::LanenetInference(const std::string& model_path) 
    : model_path_(model_path) {}

bool LanenetInference::init() {
    try {
        // 加载模型
        model_ = tflite::FlatBufferModel::BuildFromFile(model_path_.c_str());
        if (!model_) {
            throw std::runtime_error("Failed to load model");
        }

        // 创建解释器
        tflite::ops::builtin::BuiltinOpResolver resolver;
        tflite::InterpreterBuilder builder(*model_, resolver);
        builder(&interpreter_);

        if (!interpreter_) {
            throw std::runtime_error("Failed to create interpreter");
        }

        // 配置线程数和输入尺寸
        interpreter_->SetNumThreads(4);
        interpreter_->ResizeInputTensor(0, {1, 256, 512, 3});
        
        if (interpreter_->AllocateTensors() != kTfLiteOk) {
            throw std::runtime_error("Failed to allocate tensors");
        }

        return true;
    } catch (const std::exception& e) {
        std::cerr << "Initialization error: " << e.what() << std::endl;
        return false;
    }
}

std::string LanenetInference::process(const cv::Mat& input_image) {
    try {
        // 预处理图像
        cv::Mat resized_image;
        cv::resize(input_image, resized_image, cv::Size(512, 256));
        cv::Mat float_image;
        resized_image.convertTo(float_image, CV_32F, 1.0/127.5, -1.0);

        // 拷贝数据到输入张量
        float* input_data = interpreter_->typed_input_tensor<float>(0);
        memcpy(input_data, float_image.data, 256 * 512 * 3 * sizeof(float));

        // 执行推理
        if (interpreter_->Invoke() != kTfLiteOk) {
            throw std::runtime_error("Failed to invoke interpreter");
        }

        // 获取输出并处理
        TfLiteTensor* instance_tensor = interpreter_->tensor(interpreter_->outputs()[1]);
        return determine_direction(instance_tensor->data.f, 256, 512);

    } catch (const std::exception& e) {
        std::cerr << "Processing error: " << e.what() << std::endl;
        return "Error";
    }
}

bool LanenetInference::is_lane_pixel(const float* instance_seg_data, 
                                   int idx, int num_channels) {
    float sum = 0.0f;
    for (int c = 0; c < num_channels; ++c) {
        sum += instance_seg_data[idx * num_channels + c];
    }
    return (sum / num_channels) > -0.2f;
}

float LanenetInference::get_lane_center(const float* instance_seg_data,
                                      int y, int width, int num_channels) {
    int center_x = width / 2;
    int left_x = -1, right_x = -1;
    
    // 从中心向两边扫描
    for (int x = center_x; x >= 0; x--) {
        if (is_lane_pixel(instance_seg_data, y * width + x, num_channels)) {
            left_x = x;
            break;
        }
    }
    
    for (int x = center_x; x < width; x++) {
        if (is_lane_pixel(instance_seg_data, y * width + x, num_channels)) {
            right_x = x;
            break;
        }
    }
    
    // 如果找到了两边的车道线，返回中心点
    if (left_x >= 0 && right_x >= 0) {
        return (left_x + right_x) / 2.0f;
    }
    return -1;  // 未找到有效的车道线
}

std::string LanenetInference::determine_direction(const float* instance_seg_data,
                                                int height, int width) {
    const int SCAN_LINES = 10;
    const int num_channels = 4;
    std::vector<float> center_points;
    
    // 从下到上扫描获取中心点
    for(int i = 0; i < SCAN_LINES; i++) {
        int y = height - (i * height/SCAN_LINES);
        float center = get_lane_center(instance_seg_data, y, width, num_channels);
        if(center >= 0) {
            center_points.push_back(center);
        }
    }
    
    if(center_points.empty()) {
        return "Keep Straight";  // 没有检测到足够的车道线
    }
    
    // 计算近处偏移和趋势
    float near_offset = center_points[0] - width/2;
    float trend = 0;
    
    // 修复这里的语法错误
    for(size_t i = 1; i < center_points.size(); i++) {
        trend = trend + (center_points[i] - center_points[i-1]);  // 修改这行
    }
    
    if(center_points.size() > 1) {
        trend /= (center_points.size() - 1);
    }
    
    // 决策逻辑
    const float OFFSET_THRESHOLD = 10.0f;
    const float TREND_THRESHOLD = 5.0f;
    
    if(abs(near_offset) > OFFSET_THRESHOLD) {
        // 优先处理近处偏移
        return (near_offset > 0) ? "Turn Left" : "Turn Right";
    } else if(abs(trend) > TREND_THRESHOLD) {
        // 处理弯道趋势
        return (trend > 0) ? "Turn Left Slight" : "Turn Right Slight";
    } else {
        return "Keep Straight";
    }
}